
#include "select_v2_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "graph/operator_factory.h" 
#include "graph/utils/type_utils.h"
#include <algorithm>

const uint32_t BLOCK_SIZE=32;
const int32_t BUFFER_NUM = 2;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    SelectV2TilingData tiling;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());  
    auto coreNum = ascendcPlatform.GetCoreNum();  
  
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();  
    uint32_t typeBytes = 0;
    typeBytes = ge::GetSizeByDataType(context->GetInputDesc(1)->GetDataType());
    uint32_t inputLength = inputNum * typeBytes;
  
    uint64_t ubSize; 
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);   
    uint32_t uB_BlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM);
  
    ge::DataType dataType = context->GetInputTensor(1)->GetDataType();
    // uint32_t scale = GetSizeByDataType(context->GetInputDesc(0)->GetDataType())
    //                   +GetSizeByDataType(dataType)
    //                   +GetSizeByDataType(dataType)
    //                   +GetSizeByDataType(dataType);
    uint32_t scale = 8;
  
    uint32_t uselessBlock = uB_BlockNum % scale;
    uint32_t useBlock = uB_BlockNum - uselessBlock;
    uint32_t smallestVariableAllocBlock = useBlock/scale;
  
    uint32_t tileBlockNum = smallestVariableAllocBlock;
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE)/typeBytes;
  
    uint32_t inputLengthAlgin32 = ((( inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
  
    coreNum = (coreNum < inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
    coreNum = (coreNum >= 1) ? coreNum:1;
  
    uint32_t everyCoreInputBlockNum = inputLengthAlgin32/BLOCK_SIZE/coreNum;
    uint32_t tailBlockNum = inputLengthAlgin32 / BLOCK_SIZE % coreNum;
  
    uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / typeBytes;
    uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
    uint32_t smallTailDataNum = smallCoreDataNum  - (tileDataNum * smallTileNum);
    smallTailDataNum = smallTailDataNum == 0? tileDataNum : smallTailDataNum;

    //计算是否需要广播
    uint32_t tilingKey = 0;
    ge::DataType dataType1 = context->GetInputDesc(1)->GetDataType();
    if(dataType1 == ge::DT_INT32 || dataType1 == ge::DT_INT8){
        tilingKey = 3;
    }else {
        tilingKey = 2;
    }

    uint32_t conditionLength = context->GetInputTensor(0)->GetShapeSize();
    uint32_t x1Length = context->GetInputTensor(1)->GetShapeSize();
    uint32_t x2Length = context->GetInputTensor(2)->GetShapeSize();
    uint32_t yLength = context->GetOutputShape(0)->GetStorageShape().GetShapeSize();

    if(yLength!=conditionLength || yLength!=x1Length || yLength!=x2Length){
        //广播场景:输入输出形状不一致
        int32_t y_dimNum,x1_dimNum,x2_dimNum,condition_dimNum;
        int32_t y_dims[20], condition_dims[20], x1_dims[20], x2_dims[20];
        int32_t y_sumDims[20], condition_sumDims[20], x1_sumDims[20], x2_sumDims[20];

        auto shape_y = context->GetOutputShape(0)->GetOriginShape();
        auto shape_condition = context->GetInputShape(0)->GetOriginShape();
        auto shape_x1 = context->GetInputShape(1)->GetOriginShape();
        auto shape_x2 = context->GetInputShape(2)->GetOriginShape();

        //y_dimNum = shape_y.GetDimNum();
        condition_dimNum = shape_condition.GetDimNum();
        x1_dimNum = shape_x1.GetDimNum();
        x2_dimNum = shape_x2.GetDimNum();
        y_dimNum = std::max(std::max(condition_dimNum,x1_dimNum),x2_dimNum);
        //维度对齐：
        for (int i = 0; i < y_dimNum; i++)
        {
            y_dims[y_dimNum-i-1] = shape_y.GetDim(i);

            if(i < condition_dimNum){
                condition_dims[condition_dimNum-i-1] = shape_condition.GetDim(i);
            }else{
                condition_dims[i] = 1;
            }

            if(i < x1_dimNum){
                x1_dims[x1_dimNum-i-1] = shape_x1.GetDim(i);
            }else{
                x1_dims[i] = 1;
            }

            if(i < x2_dimNum){
                x2_dims[x2_dimNum-i-1] = shape_x2.GetDim(i);
            }else{
                x2_dims[i] = 1;
            }
        }

        //计算累乘维度
        y_sumDims[0] = 1;
        condition_sumDims[0] = 1;
        x1_sumDims[0] = 1;
        x2_sumDims[0] = 1;
        for(int i =1; i <= y_dimNum; i++){
            y_sumDims[i] = y_sumDims[i-1] * y_dims[i-1];
            condition_sumDims[i] = condition_sumDims[i-1] * condition_dims[i-1];
            x1_sumDims[i] = x1_sumDims[i-1] * x1_dims[i-1];
            x2_sumDims[i] = x2_sumDims[i-1] * x2_dims[i-1];
        }

        tiling.BCT.set_y_dimNum(y_dimNum);
        tiling.BCT.set_y_dims(y_dims);
        tiling.BCT.set_condition_dims(condition_dims);
        tiling.BCT.set_x1_dims(x1_dims);
        tiling.BCT.set_x2_dims(x2_dims);
        tiling.BCT.set_y_sumDims(y_sumDims);
        tiling.BCT.set_condition_sumDims(condition_sumDims);
        tiling.BCT.set_x1_sumDims(x1_sumDims);
        tiling.BCT.set_x2_sumDims(x2_sumDims);

        tilingKey = 1;

    }

    tiling.set_size(inputNum);
    tiling.set_smallCoreDataNum(smallCoreDataNum);
    tiling.set_finalSmallTileNum(finalSmallTileNum);
    tiling.set_tileDataNum(tileDataNum);
    tiling.set_smallTailDataNum(smallTailDataNum);
    tiling.set_tailBlockNum(tailBlockNum);
    tiling.set_tilingKey(tilingKey);
  
    context->SetBlockDim(coreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), 
                         context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
  
    currentWorkspace[0] = 0;
    
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{

    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
    // auto condition_dim = context->GetInputShape(0)->GetDimNum();
    // auto x1_dim = context->GetInputShape(1)->GetDimNum();
    // auto x2_dim = context->GetInputShape(2)->GetDimNum();

    // int dim = std::max(std::max(condition_dim,x1_dim),x2_dim);
    // printf("dim:%d",dim);
    // gert:: Shape* y_shape = context->GetOutputShape(0);

    // if(x1_dim == dim){
    //     const gert:: Shape* x1_shape = context->GetInputShape(1);  
    //     *y_shape = *x1_shape;

    // }else if(x2_dim == dim){
    //     const gert:: Shape* x2_shape = context->GetInputShape(2);  
    //     *y_shape = *x2_shape;
    // }else{
    //     const gert:: Shape* condition_shape = context->GetInputShape(0);  
    //     *y_shape = *condition_shape;
    // }
    return GRAPH_SUCCESS;
}
}


namespace ops {
class SelectV2 : public OpDef {
public:
    explicit SelectV2(const char* name) : OpDef(name)
    {
        this->Input("condition")
            .ParamType(REQUIRED)
            .DataType({ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL, ge::DT_BOOL})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(SelectV2);
}
